package com.rds.benchmark.connector;

import com.rds.benchmark.config.ArgvsConfig;
import io.lettuce.core.RedisURI;
import io.lettuce.core.cluster.ClusterClientOptions;
import io.lettuce.core.cluster.ClusterTopologyRefreshOptions;
import io.lettuce.core.cluster.RedisClusterClient;
import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
import io.lettuce.core.support.ConnectionPoolSupport;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;

import java.time.Duration;

public class ClusterLettuceConnector implements JedisConnector {

    private GenericObjectPool<StatefulRedisClusterConnection<String, String>> pool;
    private String stringDes;

    public ClusterLettuceConnector() {
        RedisURI redisUri = null;
        if (ArgvsConfig.password == null)
            redisUri = RedisURI.Builder.redis(ArgvsConfig.host, ArgvsConfig.port)
                    .build();
        else
            redisUri = RedisURI.Builder.redis(ArgvsConfig.host, ArgvsConfig.port)
                    .withPassword(ArgvsConfig.password)  // 添加密码
                    .build();


        ClusterTopologyRefreshOptions refreshOptions = ClusterTopologyRefreshOptions.builder()
                .enablePeriodicRefresh(true)
                .enableAllAdaptiveRefreshTriggers()
                .enablePeriodicRefresh(Duration.ofSeconds(3))
                .build();

        ClusterClientOptions options = ClusterClientOptions.builder()
                .nodeFilter(node -> !node.getFlags().contains(io.lettuce.core.cluster.models.partitions.RedisClusterNode.NodeFlag.FAIL))
                .topologyRefreshOptions(refreshOptions)
                .build();

        RedisClusterClient clusterClient = RedisClusterClient.create(redisUri);
        clusterClient.setOptions(options);

//        RedisClusterClient clusterClient = RedisClusterClient.create(redisUri);



        GenericObjectPoolConfig objectGenericObjectPoolConfig = new GenericObjectPoolConfig<>();
        objectGenericObjectPoolConfig.setMaxTotal(ArgvsConfig.threads + 1);
        objectGenericObjectPoolConfig.setMaxIdle(ArgvsConfig.threads + 1);

        pool = ConnectionPoolSupport
                .createGenericObjectPool(() -> {
                    StatefulRedisClusterConnection<String, String> connect = clusterClient.connect();
                    connect.setTimeout(Duration.ofSeconds(3));
                    return connect;
                }, objectGenericObjectPoolConfig);

        stringDes = "ClusterLettuceConnector(host='" + ArgvsConfig.host + "', port='" + ArgvsConfig.port + "' )";
        System.out.println(stringDes + " connecting....");
    }

    @Override
    public Object getJedis() {
        try {
            return pool.borrowObject();
        } catch (Exception e) {

        }
        return null;
    }

    @Override
    public boolean test() {
        try (StatefulRedisClusterConnection<String, String> connection = pool.borrowObject()) {
            RedisAdvancedClusterCommands<String, String> commands = connection.sync();
            return "PONG".equalsIgnoreCase(commands.ping());
        } catch (Exception e) {

        }
        return false;
    }

    @Override
    public String type() {
        return "CLUSTER";
    }

    @Override
    public void close(Object o){
        if(o instanceof RedisAdvancedClusterCommands){
            StatefulRedisClusterConnection statefulConnection = ((RedisAdvancedClusterCommands) o).getStatefulConnection();
            pool.returnObject(statefulConnection);
        }
    }
}
